package com.bigdata.hive;


import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;

import javax.security.auth.message.AuthException;
import java.sql.*;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * Hive 数据库操作
 *
 * @author xin.ding
 * @date 2020/10/16 13:51
 */
public class HiveDao {
    HiveConfig config;

    /**
     * Hive 连接
     */
    private Connection conn;

    public HiveDao(HiveConfig config) {
        this.config = config;
        initConnection();
    }


    /**
     * 获取Hive Jdbc连接
     *
     * @return Connection
     */
    public Connection getConnection() {
        try {
            if (conn == null || conn.isClosed()) {
                initConnection();
            }
        } catch (Exception e) {
            throw new HiveExecutorException(e);
        }
        return conn;
    }

    /**
     * 关闭连接
     */
    public void closeConnection() {
        close(conn);
    }

    /**
     * 创建Hive Jdbc连接
     */
    private void initConnection() {
        try {
            if (config.isKerberos()) {
                buildHiveKerberosEnv(config.getKrb5File(), config.getPrincipal(), config.getKeytabFile());
            }
            Class.forName("org.apache.hive.jdbc.HiveDriver");
            conn = DriverManager.getConnection(config.getUrl(), config.getUsername(), config.getPassword());
        } catch (Exception e) {
            throw new HiveExecutorException(e);
        }
    }

    /**
     * 构建HiveKerberos 认证环境
     *
     * @param krb5Path   krb5.conf 路径
     * @param principal  kerberos principal
     * @param keytabPath keytab文件路径
     */
    public static void buildHiveKerberosEnv(String krb5Path, String principal, String keytabPath) {
        System.setProperty("java.security.krb5.conf", krb5Path);
        Configuration conf = new Configuration();
        conf.set("hadoop.security.authentication", "Kerberos");
        UserGroupInformation.setConfiguration(conf);
        try {
            UserGroupInformation.loginUserFromKeytab(principal, keytabPath);
        } catch (Exception e) {
            throw new HiveExecutorException(e);
        }
    }

    /**
     * sql执行
     *
     * @param sql sql语句
     * @return 影响行数
     */

    public void execute(String sql) {
        try (Statement stmt = getConnection().createStatement()) {
            stmt.executeUpdate(sql);
        } catch (SQLException e) {
            throw new HiveExecutorException(e);
        }
    }


    /**
     * 关闭所有
     *
     * @param autoCloseables 需关闭对象
     */
    public void close(AutoCloseable... autoCloseables) {
        for (AutoCloseable autoCloseable : autoCloseables) {
            if (autoCloseable != null) {
                try {
                    autoCloseable.close();
                } catch (Exception e) {
                    throw new HiveExecutorException(e);
                }
            }
        }
    }

    /**
     * 列出指定数据库下的所有表
     */
    public List<String> showTables(String dbName) {
        Statement stmt = null;
        ResultSet res = null;
        List<String> tables = new LinkedList<>();
        try {
            stmt = getConnection().createStatement();
            if (dbName != null && dbName.trim().length() > 0) {
                stmt.execute(String.format("USE %s", dbName));
            }
            res = stmt.executeQuery("SHOW TABLES");
            while (res.next()) {
                tables.add(res.getString(1));
            }
        } catch (SQLException e) {
            throw new HiveExecutorException(e);
        } finally {
            close(res, stmt);
        }
        return tables;
    }

    /**
     * 查看数据库
     */
    public List<String> showDatabases() {
        Statement stmt = null;
        ResultSet res = null;
        List<String> tables = new LinkedList<>();
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("SHOW DATABASES");
            while (res.next()) {
                tables.add(res.getString(1));
            }
        } catch (SQLException e) {
            throw new HiveExecutorException(e);
        } finally {
            close(res, stmt);
        }
        return tables;
    }


    /**
     * 使用Statement查询数据，返回List集合
     *
     * @param sql sql
     * @return List
     */
    public List<Map<String, Object>> queryForList(String sql) {
        Statement stmt = null;
        ResultSet res = null;
        List<Map<String, Object>> list = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery(sql);
            Map<String, Object> map;
            ResultSetMetaData rsmd = res.getMetaData();
            int rowCnt = rsmd.getColumnCount();
            list = new LinkedList<>();
            while (res.next()) {
                map = new LinkedHashMap<>(rowCnt);
                for (int i = 1; i <= rowCnt; i++) {
                    map.put(rsmd.getColumnName(i), res.getObject(i));
                }
                list.add(map);
            }
        } catch (SQLException e) {
            throw new HiveExecutorException(e);
        } finally {
            close(res, stmt);
        }
        return list;
    }
}


